import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat


def normalize(u):
    if np.linalg.norm(u) == 0:
        return u
    return u / np.linalg.norm(u)


def K_Means(dataset):
    # initialize K cluster centers
    mu = np.array([np.mean(dataset[:100], axis=0), np.mean(dataset[100:], axis=0)])

    # iterate
    while True:
        # step 1: assign each sample to the nearest cluster center
        assign = np.array([np.argmin(np.linalg.norm(mu - x, axis=1)) for x in dataset])

        # step 2: update cluster center
        mu_new = np.array([np.sum(dataset[assign == i], axis=0) / np.sum(assign == i) for i in range(2)])

        # step 3: if coverage, end the iteration
        if np.linalg.norm(mu_new - mu) < 0.000001:
            return assign
        mu = mu_new


def spectral_clustering_ng(dataset, sigma, k):
    W = np.array([[np.exp(-np.linalg.norm(dataset[i] - dataset[j]) ** 2 / 2 / sigma ** 2)
                   for j in range(dataset.shape[0])] for i in range(dataset.shape[0])])
    for i in range(dataset.shape[0]):
        W[i][i] = 0
        W[i][np.argsort(W[i])[:-k]] = 0
    W = (W.T + W) / 2
    D = np.diag([np.sum(W[i]) ** (-1 / 2) for i in range(dataset.shape[0])])
    L = np.eye(dataset.shape[0]) - D.dot(W).dot(D)

    U, S, V = np.linalg.svd(L)
    u = np.real(U[:, -k:])
    u = np.array([normalize(u[i]) for i in range(dataset.shape[0])])

    assign = K_Means(u)
    return plot_cluster(dataset, assign)


def plot_cluster(dat, assign):
    # for i in range(2):
    #     plt.scatter(dat[assign == i][:, 0], dat[assign == i][:, 1])
    cl_id = [np.argmax(np.bincount(assign[i * 100:i * 100 + 100])) for i in range(2)]
    accu = np.sum([np.sum(assign[i * 100:i * 100 + 100] == cl_id[i]) for i in range(2)]) / dat.shape[0]
    # plt.title('Accu = {}'.format(accu))
    # plt.show()
    return accu


if __name__ == '__main__':
    data = np.array(loadmat('data/X_SC.mat')['X_SC'])
    Sigma = [0.01, 0.05, 0.1, 0.5, 1, 5, 10, 50, 100]
    K = [2, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90, 95, 100]
    acc = np.zeros((9, 21))
    for a in range(9):
        for b in range(21):
            acc[a][b] = spectral_clustering_ng(data, Sigma[a], K[b])
            print('SC {},{} is done'.format(a, b))

    print(acc)
    for p in range(9):
        plt.plot(K, acc[p, :], label='Sigma={}'.format(Sigma[p]))
    plt.legend()
    plt.title('Result of Spectral Clustering')
